import time
from copy import deepcopy

import path
import pandas as pd

import gxl_ai_utils.thread.my_thread
from gxl_ai_utils.utils import utils_file
from nisqa.NISQA_model import nisqaModel
import argparse


parser = argparse.ArgumentParser()
parser.add_argument('--mode', default='predict_dir', type=str,
                    help='either predict_file, predict_dir, or predict_csv')
parser.add_argument('--pretrained_model', default='weights/nisqa_tts.tar', type=str,
                    help='file name of pretrained model (must be in current working folder)')
parser.add_argument('--deg', default='./output/wav/dev/BAC009S0724W0121.wav', type=str, help='path to speech file')
parser.add_argument('--data_dir', default='./output/wav/', type=str, help='folder with speech files')
parser.add_argument('--output_dir', default='./output/', type=str, help='folder to ouput results.csv')
parser.add_argument('--csv_file', type=str, help='file name of csv (must be in current working folder)')
parser.add_argument('--csv_deg', type=str, help='column in csv with files name/path')
parser.add_argument('--num_workers', type=int, default=0, help='number of workers for pytorchs dataloader')
parser.add_argument('--bs', type=int, default=1, help='batch size for predicting')
parser.add_argument('--ms_channel', default=1, type=int, help='audio channel in case of stereo file')
parser.add_argument('--tr_device', default='cpu', type=str, )

args = parser.parse_args()
args = vars(args)

if args['mode'] == 'predict_file':
    if args['deg'] is None:
        raise ValueError('--deg argument with path to input file needed')
elif args['mode'] == 'predict_dir':
    if args['data_dir'] is None:
        raise ValueError('--data_dir argument with folder with input files needed')
elif args['mode'] == 'predict_csv':
    if args['csv_file'] is None:
        raise ValueError('--csv_file argument with csv file name needed')
    if args['csv_deg'] is None:
        raise ValueError('--csv_deg argument with csv column name of the filenames needed')
    if args['data_dir'] is None:
        args['data_dir'] = ''
else:
    raise NotImplementedError('--mode given not available')
args['tr_bs_val'] = args['bs']
args['tr_num_workers'] = args['num_workers']


def get_qisqa_mos_from_wav_dir(wav_dir: str, output_dir: str):
    print(f'开始处理如下目录中的音频: {wav_dir}')
    args_my = deepcopy(args)
    args_my['data_dir'] = wav_dir
    utils_file.makedir_sil(output_dir)
    args_my['output_dir'] = output_dir
    nisqa = nisqaModel(args_my)
    nisqa.predict()


def read_csv_to_dic(csv_path: str):
    df = pd.read_csv(csv_path)
    deg_list = df['deg'].to_list()
    mos_pred_list = df['mos_pred'].to_list()
    print(deg_list)
    print(mos_pred_list)
    res_dic = {}
    for k, v in zip(deg_list, mos_pred_list):
        ks = k.strip().split('_')
        if len(ks) >= 2:
            k = k.strip().split('_')[0] + "_" + k.strip().split('_')[1]
        res_dic[k] = v
    return res_dic


def do_for_one_dir(dirpath: str, output_dir: str, res_dict: dict):
    get_qisqa_mos_from_wav_dir(dirpath, output_dir)
    dic = read_csv_to_dic(utils_file.join_path(output_dir, 'NISQA_results.csv'))
    res_dict.update(dic)


def do_multi_thread():
    thread_pool = gxl_ai_utils.thread.my_thread.MyThreadPool()
    dir_path_list = ['./output/wav','./output/wav2']
    output_dir = './output'
    res_dict = {}
    for i, dir_path in enumerate(dir_path_list):
        thread_pool.add_thread(do_for_one_dir, [dir_path, utils_file.join_path(output_dir, f'dir_{i}'), res_dict])
    thread_pool.start()
    utils_file.write_dic_to_scp(res_dict, utils_file.join_path(output_dir, 'NISQA_mos.scp'))


if __name__ == "__main__":
    """"""
    now = time.time()
    little_aishell_path = 'E:\gengxuelong_study\server_local_adapter\\ai\data\small_aishell\dev'
    # do_for_one_dir(little_aishell_path, './output/12')
    do_multi_thread()
    end = time.time()
    print(f'耗时{end - now}s')
